#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: liang kang
@contact: gangkanli1219@gmail.com
@time: 2018/5/30 16:38
@desc: 
"""
from pathlib import Path
from queue import Queue

import h5py
import numpy as np
from skimage import measure, morphology, segmentation

from utils.config import DataPipe as DP
from utils.config import Number
from utils.data import Image, Seed
from utils.functions import average_distance2point
from utils.log import get_console_logger


def _get_region(array, mask, point, elem):
    shape = np.asarray(elem.shape)
    center = np.asarray(point)
    start = center - shape // 2
    end = start + shape
    new_arr = array[start[0]:end[0], start[1]:end[1], start[2]:end[2]]
    new_arr = np.logical_and(elem, new_arr)
    new_mask = mask[start[0]:end[0], start[1]:end[1], start[2]:end[2]]
    return new_arr, new_mask, start


def _set_region(array, sub_arr, center):
    shape = np.asarray(sub_arr.shape)
    center = np.asarray(center)
    start = center - shape // 2
    end = start + shape
    array[start[0]:end[0], start[1]:end[1], start[2]:end[2]] = sub_arr
    return array


class Airway(object):
    __slots__ = ('msg',
                 'mask',
                 'head',
                 'seed',
                 'point_queue',
                 'save',
                 'logger')

    def __init__(self, mask: Image, save: str=None, logger=None):
        self.msg = None
        self.mask = mask
        self.save = save
        if logger is not None:
            self.logger = logger
        else:
            self.logger = get_console_logger('Airway')

    def _head_direct(self):
        """
        find the direct of head

        Returns
        -------

        """
        image_array = self.mask.image_array
        shape = image_array.shape
        upside = image_array[0:Number.HEAD_RANGE, :, :]
        up_center = (Number.HEAD_RANGE // 2, shape[1] // 2, shape[2] // 2)
        upside_avg_dis = average_distance2point(upside, up_center)
        downside = image_array[(shape[2] - Number.HEAD_RANGE):, :, :]
        down_center = (shape[2] - Number.HEAD_RANGE // 2,
                       shape[1] // 2, shape[2] // 2)
        downside_avg_dis = average_distance2point(downside, down_center)
        self.head = upside_avg_dis > downside_avg_dis

    def _find_airway(self):
        image_array = self.mask.image_array.copy()
        if not self.head:
            image_array = np.flip(image_array, axis=0)

        upside = image_array[0:Number.HEAD_RANGE, :, :]
        labeled = measure.label(upside)
        regions = measure.regionprops(labeled)
        airways = []
        for region in regions:
            box = region.bbox
            if (box[3] - box[0]) > (Number.AIRWAY_RANGE // 2):
                areas = np.sum(
                    upside[box[0]:box[3], box[1]:box[4], box[2]:box[5]],
                    axis=(1, 2))
                if np.mean(areas) < Number.AIRWAY_AREA:
                    airways.append(region)

        if len(airways) == 0:
            self.msg = 'This No Airway'
            self.seed = None
            return
        elif len(airways) > 1:
            airway = min(airways, key=lambda a: a.area)
        else:
            airway = airways[0]

        airway_mask = np.equal(airway.label, labeled)
        airway_mask = morphology.erosion(
            airway_mask, selem=morphology.ball(Number.AIRWAY_BALL + 2))
        points = np.nonzero(airway_mask)
        self.seed = np.array([points[0][0] + Number.AIRWAY_ADD_SIZE,
                              points[1][0], points[2][0]])

    def _growing(self, mask_region, image_region,
                 seed, label, start,
                 need_labeled=False):
        """
        Rule

        Returns
        -------

        """
        center = np.asarray(image_region.shape) // 2 + 1
        if need_labeled:
            labeled = measure.label(image_region)
            labeled[labeled != labeled[np.split(center, 3, 0)]] = 0
            image = labeled.astype(np.bool)
        else:
            image = image_region
        mask_region_ = mask_region.copy()
        mask_region_[np.logical_and(
            np.equal(0, mask_region), image)] = max(1, label - 1)
        mask_region_[np.split(center, 3, 0)] = label
        points = np.logical_xor(mask_region, mask_region_)
        points = segmentation.find_boundaries(points, mode='inner')
        points = np.asarray(np.nonzero(points)).T
        points = np.add(points, np.expand_dims(start, axis=0))
        if points.size != 0:
            if label == 3:
                parent = seed
            elif label == 2:
                parent = seed.parent
            else:
                parent = seed.parent.parent
            for point in points:
                self.point_queue.put(Seed(point, parent))
        return mask_region_

    def _remove_airway(self):
        """
        Region Grown

        Returns
        -------

        """
        image_array = self.mask.image_array.copy()
        if not self.head:
            image_array = np.flip(image_array, axis=0)
        image_array = np.concatenate(
            [np.zeros([Number.AIRWAY_ADD_SIZE,
                       image_array.shape[1],
                       image_array.shape[2]], dtype=np.bool),
             image_array],
            axis=0)
        airway_mask = np.zeros_like(image_array, dtype=np.uint8)

        self.point_queue = Queue()
        airway_mask[np.split(self.seed, 3, 0)] = 3
        self.point_queue.put(Seed(self.seed, Seed(self.seed, None)))

        ball_elem = morphology.ball(Number.AIRWAY_BALL)
        ball_summary = np.sum(ball_elem)
        cross_elem = morphology.octahedron(Number.AIRWAY_CROSS)

        while not self.point_queue.empty():
            seed = self.point_queue.get()
            if np.linalg.norm(np.subtract(seed.point, seed.parent.point)) >= (
                    Number.AIRWAY_BALL + Number.AIRWAY_CROSS + 1):
                continue
            if airway_mask[np.split(seed.point, 3, 0)] > 1:
                image_region, mask_region, start = _get_region(
                    image_array, airway_mask, seed.point, ball_elem)
                if ball_summary == np.sum(image_region):
                    mask_region = self._growing(
                        mask_region, image_region, seed, 3, start)
                # elif np.sum(image_region) == 0:
                #     continue
                else:
                    mask_region = self._growing(
                        mask_region, image_region, seed, 2, start,
                        need_labeled=True)
            else:
                image_region, mask_region, start = _get_region(
                    image_array, airway_mask, seed.point, cross_elem)
                # if np.sum(image_region) == 0:
                #     continue
                mask_region = self._growing(
                    mask_region, image_region, seed, 1, start,
                    need_labeled=True)

            airway_mask = _set_region(
                airway_mask, mask_region, seed.point)
        airway_mask = airway_mask[Number.AIRWAY_ADD_SIZE:, :, :]
        image_array = image_array[Number.AIRWAY_ADD_SIZE:, :, :]
        self.mask.image_array = (
            np.logical_and(np.logical_not(airway_mask), image_array))

        if self.save is not None:
            self.logger.info('Saving the data')
            path = Path(self.save) / DP.AIRWAY + '.h5'
            with h5py.File(str(path), 'w') as file:
                file.create_dataset(DP.AIRWAY_MASK, data=airway_mask,
                                    compression='gzip', compression_opts=9)
                file.create_dataset(DP.LUNG_NO_AIRWAY,
                                    data=self.mask.image_array,
                                    compression='gzip', compression_opts=9)

    def update(self):
        self.logger.info('Removing the Airway ...')
        self.logger.info('Finding the direct of head')
        self._head_direct()
        self.logger.info('Finding the seed of airway')
        self._find_airway()
        if self.msg is not None:
            return
        self.logger.info('Removing the Airway using Region Growing')
        self._remove_airway()
